"""
嵌入缓存模块

提供嵌入向量的缓存功能，提高重复查询的性能。
"""

import hashlib
import pickle
import time
from typing import Dict, Optional
from dataclasses import dataclass


@dataclass
class EmbeddingCache:
    """嵌入缓存项"""
    text_hash: str
    embedding: list
    timestamp: float
    model_name: str


class EmbeddingCacheManager:
    """嵌入缓存管理器"""

    def __init__(self, cache_ttl: int = 3600):
        """
        初始化缓存管理器

        Args:
            cache_ttl: 缓存TTL（秒）
        """
        self._cache: Dict[str, EmbeddingCache] = {}
        self._cache_ttl = cache_ttl
        self._cache_enabled = True

    def set_ttl(self, ttl: int) -> None:
        """设置缓存TTL"""
        self._cache_ttl = ttl

    def set_enabled(self, enabled: bool) -> None:
        """启用/禁用缓存"""
        self._cache_enabled = enabled

    def get(self, text: str) -> Optional[list]:
        """
        获取文本的嵌入向量

        Args:
            text: 输入文本

        Returns:
            Optional[list]: 嵌入向量，如果缓存不存在或已过期则返回None
        """
        if not self._cache_enabled:
            return None

        text_hash = self._get_text_hash(text)
        if text_hash in self._cache and self._is_cache_valid(self._cache[text_hash]):
            return self._cache[text_hash].embedding
        return None

    def put(self, text: str, embedding: list, model_name: str) -> None:
        """
        存储文本的嵌入向量到缓存

        Args:
            text: 输入文本
            embedding: 嵌入向量
            model_name: 模型名称
        """
        if not self._cache_enabled:
            return

        text_hash = self._get_text_hash(text)
        self._cache[text_hash] = EmbeddingCache(
            text_hash=text_hash,
            embedding=embedding,
            timestamp=time.time(),
            model_name=model_name
        )

    def clear(self) -> None:
        """清空所有缓存"""
        self._cache.clear()

    def get_stats(self) -> Dict[str, any]:
        """
        获取缓存统计信息

        Returns:
            Dict[str, any]: 缓存统计数据
        """
        valid_cache_count = sum(
            1 for cache_item in self._cache.values()
            if self._is_cache_valid(cache_item)
        )

        return {
            "cache_enabled": self._cache_enabled,
            "total_cache_items": len(self._cache),
            "valid_cache_items": valid_cache_count,
            "cache_ttl": self._cache_ttl,
            "cache_memory_mb": round(len(pickle.dumps(self._cache)) / (1024 * 1024), 2)
        }

    def _get_text_hash(self, text: str) -> str:
        """生成文本的哈希值"""
        return hashlib.md5(text.encode('utf-8')).hexdigest()

    def _is_cache_valid(self, cache_item: EmbeddingCache) -> bool:
        """检查缓存项是否有效"""
        return time.time() - cache_item.timestamp < self._cache_ttl